#include "trajectory_optimization/trajectory_optimization.h"
#include "transform_utils/display.h"

namespace Planner
{
    void TrajectoryOptimization::setParams(const ros::NodeHandle &nh)
    {
        nh.param("trajectory_optimization/algorithm", algorithm_, 34);
        nh.param("trajectory_optimization/opt_time", opt_time_, 0.5);
        nh.param("trajectory_optimization/opt_iteration", opt_iteration_, 200);
        nh.param("trajectory_optimization/opt_step", opt_step_, 0.05); // 算法的初始步长
        nh.param("trajectory_optimization/order", order_, 5);
        nh.param("trajectory_optimization/continus", continus_, 3);
        nh.param("trajectory_optimization/type", type_, 3);
        nh.param("trajectory_optimization/speed", speed_, 1.0);
        nh.param("trajectory_optimization/vmax", vmax_, 2.0);
        nh.param("trajectory_optimization/amax", amax_, 2.0);
        nh.param("trajectory_optimization/resolution", resolution_, 0.1);
        nh.param("trajectory_optimization/time_gap", time_gap_, 0.05);
        nh.param("trajectory_optimization/secure_dis", secure_dis_, 1.0);
        nh.param("trajectory_optimization/wt", wt_, 1.0);
        nh.param("trajectory_optimization/ws", ws_, 1.0);
        nh.param("trajectory_optimization/wc", wc_, 10.0);
        nh.param("trajectory_optimization/wdv", wdv_, 1.0);
        nh.param("trajectory_optimization/wda", wda_, 1.0);
        nh.param("trajectory_optimization/alphac", alphac_, 1.0); // 这个或许应该采用更小的值，使得障碍物间隙惩罚项在接近安全距离时急剧增大
        nh.param("trajectory_optimization/alphav", alphav_, 0.1);
        nh.param("trajectory_optimization/alphaa", alphaa_, 0.1);
        nh.param("trajectory_optimization/rc", rc_, 1.0);
        nh.param("trajectory_optimization/rv", rv_, 1.0);
        nh.param("trajectory_optimization/ra", ra_, 1.0);
        nh.param("trajectory_optimization/pos_tol", pos_tol_, 0.1);
        nh.param("trajectory_optimization/print_info", print_info_, false);
    }

    void TrajectoryOptimization::setWayPoints(const vector<Vector3d> &pts)
    {
        way_points_ = pts;
        // traj_gen_.setWayPoints(way_points_);
        segments_ = way_points_.size() - 1;
        // 初始时间分配，两种方式
        /*
        time_seq_.reserve(way_points_.size() - 1);
        for (size_t i = 0; i < way_points_.size() - 1; i++)
        {
            double distance = (way_points_[i + 1] - way_points_[i]).norm();
            // 采用平均速度分配时间
            double t = distance / speed_;
            // 采用mav_trajectory_generation中时间分配方式
            // t = 2 * distance/vmax_ * (1 + magic_fabian_constant * vmax_/amax_
            // * exp(-distance/vmax_ * 2);
            // magic_fabian_constant was determined to 6.5 in a student project ...
            time_seq_.push_back(t);
        } */
    }

    void TrajectoryOptimization::setBoundaryCondition(const Vector3d &startVel, const Vector3d &startAcc,
                                                      const Vector3d &goalVel, const Vector3d &goalAcc)
    {
        startVel_ = startVel;
        startAcc_ = startAcc;
        goalVel_ = goalVel;
        goalAcc_ = goalAcc;
    }

    void TrajectoryOptimization::setCloudMap(pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud, pcl::KdTreeFLANN<pcl::PointXYZ>::Ptr &tree)
    {
        cloud_map_ = cloud;
        kdtree_ = tree;
    }

    int TrajectoryOptimization::optimize()
    {
        // 计算初始轨迹(初始轨迹依然采用TrajectoryGenerator计算)
        traj_gen_.setParams(order_, speed_, continus_);
        traj_gen_.setWayPoints(way_points_);
        traj_gen_.setBoundaryCondition(startVel_, startAcc_, goalVel_, goalAcc_);
        traj_gen_.solve(type_);
        traj_gen_.getCoeffs(coeffs_);
        if (print_info_)
            cout << "initial trajectory generated\n";

        // traj_gen_计算初始估计后,由traj_gen_free_完成之后工作
        traj_gen_free_.setParams(order_, speed_, continus_);
        traj_gen_free_.setWayPoints(way_points_);
        traj_gen_free_.setBoundaryCondition(startVel_, startAcc_, goalVel_, goalAcc_);
        traj_gen_free_.initArgs(type_);

        // 需要优化的参数为dp及t
        vector<MatrixXd> dfree_tmp, dfix_tmp; // dp和df需要转化
        traj_gen_.getVals(R_, Rfp_, Rpp_, dfree_tmp, dfix_tmp, AiMC_);
        traj_gen_.getTimeSeq(time_seq_);
        // dfree_tmp=[v1 a1 v2 a2 .. vn-1 an-1],dfix_tmp=[p0 p1 ... pn v0 a0 vn an]
        // 需要转化为dp=[p1 v1 a1 p2 v2 a2 .. pn-1 vn-1 an-1],df=[p0 v0 a0 pn vn an]
        dp_.clear();
        df_.clear();
        for (int axis = 0; axis < 3; axis++)
        {
            MatrixXd dp_tmp(3 * segments_ - 3, 1), df_tmp(6, 1);
            for (int row = 0; row < segments_ - 1; row++)
            {
                dp_tmp(3 * row, 0) = dfix_tmp[axis](row + 1, 0);          // pos
                dp_tmp(3 * row + 1, 0) = dfree_tmp[axis](2 * row, 0);     // vel
                dp_tmp(3 * row + 2, 0) = dfree_tmp[axis](2 * row + 1, 0); // acc
            }

            df_tmp(0, 0) = dfix_tmp[axis](0, 0);
            df_tmp(1, 0) = dfix_tmp[axis](segments_ + 1, 0);
            df_tmp(2, 0) = dfix_tmp[axis](segments_ + 2, 0);
            df_tmp(3, 0) = dfix_tmp[axis](segments_, 0);
            df_tmp(4, 0) = dfix_tmp[axis](segments_ + 3, 0);
            df_tmp(5, 0) = dfix_tmp[axis](segments_ + 4, 0);

            dp_.push_back(dp_tmp);
            df_.push_back(df_tmp);
        }

        // 初始化优化器
        int num_dp = dp_[0].rows() * 3;
        int num_opt = num_dp + segments_;
        int seed = ros::Time::now().toNSec() % 65536;
        nlopt::srand(seed);
        nlopt_.reset(new nlopt::opt(nlopt::algorithm(algorithm_), num_opt));
        nlopt_->set_maxtime(opt_time_);
        nlopt_->set_maxeval(opt_iteration_);
        // 初始解
        vector<double> initial_solution, initial_step, lower_bounds, upper_bounds;
        initial_solution.reserve(num_opt);
        initial_step.reserve(num_opt);
        lower_bounds.reserve(num_opt);
        upper_bounds.reserve(num_opt);
        for (size_t i = 0; i < 3; i++)
            for (size_t j = 0; j < dp_[0].rows(); j++)
                initial_solution.push_back(dp_[i](j, 0));
        for (size_t i = 0; i < segments_; i++)
            initial_solution.push_back(time_seq_[i]);
        // 初始步长
        for (size_t i = 0; i < num_opt; i++)
            initial_step.push_back(opt_step_);
        // 设置边界
        for (size_t i = 0; i < 3; i++)
            for (size_t j = 0; j < segments_ - 1; j++)
            {
                lower_bounds.push_back(dp_[i](3 * j, 0) - pos_tol_);
                lower_bounds.push_back(-vmax_);
                lower_bounds.push_back(-amax_);
                upper_bounds.push_back(dp_[i](3 * j, 0) + pos_tol_);
                upper_bounds.push_back(vmax_);
                upper_bounds.push_back(amax_);
            }
        for (size_t i = 0; i < segments_; i++)
        {
            double distance = (way_points_[i + 1] - way_points_[i]).norm();
            lower_bounds.push_back(distance / vmax_);
            upper_bounds.push_back(std::numeric_limits<double>::max());
        }

        try
        {
            // cout << lower_bounds.size() << endl;
            // cout << upper_bounds.size() << endl;
            // cout << initial_step.size() << endl;
            nlopt_->set_lower_bounds(lower_bounds);
            nlopt_->set_upper_bounds(upper_bounds);
            nlopt_->set_initial_step(initial_step);
            nlopt_->set_min_objective(&TrajectoryOptimization::costFunc, this);
        }
        catch (const std::exception &e)
        {
            std::cerr << "error while setting up nlopt: " << e.what() << '\n';
            return nlopt::FAILURE;
        }
        if (print_info_)
        {
            cout << "nlopt initialized\n";
            printVector(initial_solution);
        }

        double final_cost;
        int result;
        iter_num_ = 0;
        try
        {
            result = nlopt_->optimize(initial_solution, final_cost);
        }
        catch (const std::exception &e)
        {
            std::cerr << "error while running nlopt" << e.what() << '\n';
            return nlopt::FAILURE;
        }

        // 提取优化后的结果
        dp_.clear();
        time_seq_.clear();
        MatrixXd dpx(3 * segments_ - 3, 1), dpy(3 * segments_ - 3, 1), dpz(3 * segments_ - 3, 1);
        for (size_t i = 0; i < dpx.rows(); i++)
        {
            dpx(i, 0) = initial_solution[i];
            dpy(i, 0) = initial_solution[i + 3 * (segments_ - 1)];
            dpz(i, 0) = initial_solution[i + 6 * (segments_ - 1)];
        }
        dp_ = {dpx, dpy, dpz};
        for (size_t i = 0; i < segments_; i++)
            time_seq_.push_back(initial_solution[i + 9 * (segments_ - 1)]);
        // 计算coeffs和代价函数
        double cost;
        traj_gen_free_.solveByDpAndTime(dp_, time_seq_, coeffs_, cost);
        // printVector(dp_);
        if (print_info_)
            printVector(time_seq_);
        return result;
    }

    // 选取gradient-free的算法，无需计算梯度
    double TrajectoryOptimization::costFunc(const vector<double> &x, vector<double> &grad, void *func_data)
    {
        TrajectoryOptimization *opt = reinterpret_cast<TrajectoryOptimization *>(func_data);
        return opt->getCostFunc(x);
    }

    double TrajectoryOptimization::getCostFunc(const vector<double> &x)
    {
        iter_num_++;
        double cost_time, cost_smooth, cost_collision, cost_dynamic_vel, cost_dynamic_acc, cost;
        // 提取dp和time_seq
        dp_.clear();
        time_seq_.clear();
        MatrixXd dpx(3 * (segments_ - 1), 1), dpy(3 * (segments_ - 1), 1), dpz(3 * (segments_ - 1), 1);
        for (size_t i = 0; i < dpx.rows(); i++)
        {
            dpx(i, 0) = x[i];
            dpy(i, 0) = x[i + 3 * (segments_ - 1)];
            dpz(i, 0) = x[i + 6 * (segments_ - 1)];
        }
        dp_ = {dpx, dpy, dpz};
        for (size_t i = 0; i < segments_; i++)
            time_seq_.push_back(x[i + 9 * (segments_ - 1)]);
        // 计算coeffs和代价函数
        traj_gen_free_.solveByDpAndTime(dp_, time_seq_, coeffs_, cost_smooth);
        cost_time = getTimeCostFunc();

        cost_collision = cost_dynamic_vel = cost_dynamic_acc = 0;
        getCollisionCostAndDynamicCostFunc(cost_collision, cost_dynamic_vel, cost_dynamic_acc);

        cost = wt_ * cost_time + ws_ * cost_smooth + wc_ * cost_collision + wdv_ * cost_dynamic_vel + wda_ * cost_dynamic_acc;
        if (print_info_)
            cout << "iteration " << iter_num_ << ": total_cost " << cost << ", cost_smooth " << cost_smooth << ", cost_time " << cost_time
                 << ", cost_collision " << cost_collision << ", cost_vel " << cost_dynamic_vel << ", cost_acc " << cost_dynamic_acc << endl;
        cost_val_ = {cost_time, cost_smooth, cost_collision, cost_dynamic_vel, cost_dynamic_acc};
        return cost;
    }

    double TrajectoryOptimization::getTimeCostFunc()
    {
        double total_time = 0;
        for (size_t i = 0; i < time_seq_.size(); i++)
            total_time += time_seq_[i];
        return total_time * total_time;
    }

    double TrajectoryOptimization::getDistance(double x, double y, double z)
    {
        pcl::PointXYZ pt(x, y, z);
        std::vector<int> pointIdxNKNSearch(1);
        std::vector<float> pointNKNSquaredDistance(1);
        if (kdtree_->nearestKSearch(pt, 1, pointIdxNKNSearch, pointNKNSquaredDistance) > 0)
        {
            float squaredDis = pointNKNSquaredDistance[0];
            return sqrt(squaredDis);
        }
        std::cout << "kdtree search error\n";
        return 0;
    }

    double TrajectoryOptimization::getDistanceAndGrad(double x, double y, double z, Vector3d &grad)
    {
        grad[0] = (getDistance(x + resolution_, y, z) - getDistance(x - resolution_, y, z)) / (2.0 * resolution_);
        grad[1] = (getDistance(x, y + resolution_, z) - getDistance(x, y - resolution_, z)) / (2.0 * resolution_);
        grad[2] = (getDistance(x, y, z + resolution_) - getDistance(x, y, z - resolution_)) / (2.0 * resolution_);
        return getDistance(x, y, z);
    }

    double TrajectoryOptimization::getCollisionPenaltyByLoc(State &loc)
    {
        double dis = getDistance(loc.pos[0], loc.pos[1], loc.pos[2]);
        double collision_penalty = alphac_ * exp(-(dis - secure_dis_) / rc_);
        return collision_penalty * loc.vel.norm() * time_gap_;
    }

    double TrajectoryOptimization::getVelPenaltyByLoc(State &loc)
    {
        double collision_vel = alphav_ * exp(-(loc.vel.norm() - vmax_) / rv_);
        return collision_vel * loc.vel.norm() * time_gap_;
    }

    double TrajectoryOptimization::getAccPenaltyByLoc(State &loc)
    {
        double collision_acc = alphaa_ * exp(-(loc.acc.norm() - amax_) / ra_);
        return collision_acc * loc.vel.norm() * time_gap_;
    }

    void TrajectoryOptimization::getCollisionCostAndDynamicCostFunc(double &cost_collision, double &cost_dynamic_vel, double &cost_dynamic_acc)
    {
        if (wc_ < 1e-3 && wdv_ < 1e-3 && wda_ < 1e-3)
            return;

        State location;
        for (size_t i = 0; i < time_seq_.size(); i++)
        {
            // NOTE 这里可以固定分为15段,也可以固定时间间隔
            double time_gap = time_seq_[i] / 15.0;
            // FIXME 之前没有更新time_gap_的值
            time_gap_ = time_gap;
            for (double t = 1e-3; t < time_seq_[i]; t += time_gap)
            {
                traj_gen_free_.getState(t, location, i, true);
                if (wc_ > 1e-3)
                    cost_collision += getCollisionPenaltyByLoc(location);
                if (wdv_ > 1e-3)
                    cost_dynamic_vel += getVelPenaltyByLoc(location);
                if (wda_ > 1e-3)
                    cost_dynamic_acc += getAccPenaltyByLoc(location);
            }
        }
    }

} // namespace Planner